package org.jxkj.app

import org.apache.spark.sql.Row
import org.apache.spark.sql.expressions.{MutableAggregationBuffer, UserDefinedAggregateFunction}
import org.apache.spark.sql.types._
class MyUDAF extends UserDefinedAggregateFunction{
  //输入的数据类型的Schema
  override def inputSchema: StructType = StructType(Array(StructField("类型1", DoubleType)))
  /**
    * 初始化操作
    * @param buffer
    */
  override def initialize(buffer: MutableAggregationBuffer) = {
    //也可以做一些初始化安全判断的操作
    assert(buffer.length==2)
    //MutableAggregationBuffer初始化接受的两个参数，这个根据你的需要来决定参数的个数
    //最多可以延长22个，可以看看MutableAggregationBuffer的的实现
    buffer(0) = 0.0
    buffer(1) = 0L  }
  /**    * 初始化之后，每次迭代的计算    * 思想：让每次数据的Row和我们初始化的MutableAggregationBuffer做一些我们设计好的操作
    * @param buffer 我们初始化的MutableAggregationBuffer，是用来接受数值的
    * @param input 每次迭代输入的ROw
    */
  override def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
    //这个input就是我们调用自定义函数输入的参数
    buffer(0) = buffer.getDouble(0)+input.getDouble(0)
    buffer(1) = buffer.getLong(1)+1
  }
  //中间计算时候的模式
  override def bufferSchema: StructType = StructType(Array(
    //根据你update中的数据类型，来设置
    StructField("sum",DoubleType),   StructField("count",LongType)  ))
  //局部聚合操作
  /**
    * 局部聚合操作，这部分就相当于map-reduce中的combine,有map-reduce程序中的设置：job.setCombinerClass（Some.class）
    * @param buffer1 局部buffer1
    * @param buffer2 局部buffer2
    */
  override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {
    //buffer1和buffer2都是update函数中的结果，操作的时候根据需求才操作，因为我们的操作是单纯的叠加
    buffer1(0) = buffer1.getDouble(0) + buffer2.getDouble(0)
    buffer1(1) = buffer1.getLong(1) + buffer2.getLong(1)
  }
  /**
    * Returns true iff this function is deterministic, i.e. given the same input,
    * always return the same output.
    */
  override def deterministic: Boolean = true
  /**
    * 最终平均值
    * @param buffer 就是我们merge函数的中结果
    * @return 我们想要的数值
    */
  override def evaluate(buffer: Row): Any = {
    buffer.getDouble(0)/buffer.getLong(1)
  }
  //返回的数据类型
  override def dataType: DataType = DoubleType
}
//spark.udf.register("average", new AverageSal)